{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "import numpy as np\n",
    "import core.utils as utils\n",
    "import tensorflow as tf\n",
    "import os\n",
    "from core.yolov3 import YOLOv3, decode\n",
    "from PIL import Image\n",
    "\n",
    "os.environ['TF_CPP_MIN_LOG_LEVEL']='3'\n",
    "\n",
    "input_size   = 416\n",
    "image_path   = \"./docs/road.jpg\"\n",
    "\n",
    "input_layer  = tf.keras.layers.Input([input_size, input_size, 3])\n",
    "feature_maps = YOLOv3(input_layer)\n",
    "\n",
    "original_image      = cv2.imread(image_path)\n",
    "original_image      = cv2.cvtColor(original_image, cv2.COLOR_BGR2RGB)\n",
    "original_image_size = original_image.shape[:2]\n",
    "\n",
    "image_data = utils.image_preporcess(np.copy(original_image), [input_size, input_size])\n",
    "image_data = image_data[np.newaxis, ...].astype(np.float32)\n",
    "\n",
    "bbox_tensors = []\n",
    "for i, fm in enumerate(feature_maps):\n",
    "    bbox_tensor = decode(fm, i)\n",
    "    bbox_tensors.append(bbox_tensor)\n",
    "\n",
    "model = tf.keras.Model(input_layer, bbox_tensors)\n",
    "utils.load_weights(model, \"./yolov3.weights\")\n",
    "# model.summary()\n",
    "\n",
    "pred_bbox = model.predict(image_data)\n",
    "pred_bbox = [tf.reshape(x, (-1, tf.shape(x)[-1])) for x in pred_bbox]\n",
    "pred_bbox = tf.concat(pred_bbox, axis=0)\n",
    "bboxes = utils.postprocess_boxes(pred_bbox, original_image_size, input_size, 0.3)\n",
    "bboxes = utils.nms(bboxes, 0.45, method='nms')\n",
    "print(bboxes)\n",
    "\n",
    "image = utils.draw_bbox(original_image, bboxes)\n",
    "image = Image.fromarray(image)\n",
    "image.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
